#include "hip/hip_runtime.h"
#if CUDA_VERSION_MAJOR >= 8

#include <cassert>

#include <hip/hip_runtime.h>

#include "scope/init/init.hpp"
#include "scope/utils/utils.hpp"
#include "scope/init/flags.hpp"

#include "args.hpp"

#define NAME "Comm/UM/Latency/GPUToGPU"

template <bool NOOP = false>
__global__ void gpu_traverse(size_t *ptr, const size_t steps) {

  if (NOOP) {
    return;
  }
  size_t next = 0;
  for (int i = 0; i < steps; ++i) {
    next = ptr[next];
  }
  ptr[next] = 1;
}

static void Comm_UM_Latency_GPUToGPU(benchmark::State &state) {

  if (!has_cuda) {
    state.SkipWithError(NAME " no CUDA device found");
    return;
  }

  if (num_gpus() < 2) {
    state.SkipWithError(NAME " requires at least 2 gpus");
    return;
  }
  assert(FLAG(cuda_device_ids).size() >= 2);
  const int src_id   = FLAG(cuda_device_ids)[0];
  const int dst_id   = FLAG(cuda_device_ids)[1];

  if (src_id == dst_id) {
    state.SkipWithError(NAME "src and dst GPU should be different");
    return;
  }

  const size_t steps = state.range(0);

  const size_t stride = 65536 * 2;
  const size_t bytes  = sizeof(size_t) * (steps + 1) * stride;

  if (PRINT_IF_ERROR(utils::cuda_reset_device(src_id))) {
    state.SkipWithError(NAME " failed to reset src device");
    return;
  }
  if (PRINT_IF_ERROR(utils::cuda_reset_device(src_id))) {
    state.SkipWithError(NAME " failed to reset dst device");
    return;
  }

  if (PRINT_IF_ERROR(hipSetDevice(dst_id))) {
    state.SkipWithError(NAME " failed to set CUDA dst device");
    return;
  }

  size_t *ptr = nullptr;
  if (PRINT_IF_ERROR(hipMallocManaged(&ptr, bytes))) {
    state.SkipWithError(NAME " failed to perform hipMallocManaged");
    return;
  }
  defer(hipFree(ptr));

  if (PRINT_IF_ERROR(hipMemset(ptr, 0, bytes))) {
    state.SkipWithError(NAME " failed to perform hipMemset");
    return;
  }
  // set up stride pattern
  for (size_t i = 0; i < steps; ++i) {
    ptr[i * stride] = (i + 1) * stride;
  }
  if (PRINT_IF_ERROR(hipSetDevice(src_id))) {
    state.SkipWithError(NAME " failed to set CUDA src device");
    return;
  }
  if (PRINT_IF_ERROR(hipDeviceSynchronize())) {
    state.SkipWithError(NAME " failed to synchronize");
    return;
  }

  if (PRINT_IF_ERROR(hipSetDevice(dst_id))) {
    state.SkipWithError(NAME " failed to set CUDA dst device");
    return;
  }
  hipEvent_t start, stop;
  if (PRINT_IF_ERROR(hipEventCreate(&start))) {
    state.SkipWithError(NAME " failed to create start event");
    return;
  }
  defer(hipEventDestroy(start));

  if (PRINT_IF_ERROR(hipEventCreate(&stop))) {
    state.SkipWithError(NAME " failed to create end event");
    return;
  }
  defer(hipEventDestroy(stop));

  for (auto _ : state) {
    if (PRINT_IF_ERROR(cudaMemPrefetchAsync(ptr, bytes, src_id))) {
      state.SkipWithError(NAME " failed to prefetch to src");
      return;
    }
    if (PRINT_IF_ERROR(hipSetDevice(src_id))) {
      state.SkipWithError(NAME " failed to set CUDA src device");
      return;
    }
    if (PRINT_IF_ERROR(hipDeviceSynchronize())) {
      state.SkipWithError(NAME " failed to synchronize");
      return;
    }
    if (PRINT_IF_ERROR(hipSetDevice(dst_id))) {
      state.SkipWithError(NAME " failed to set dst device");
      return;
    }
    if (PRINT_IF_ERROR(hipDeviceSynchronize())) {
      state.SkipWithError(NAME " failed to synchronize");
      return;
    }

    hipEventRecord(start);
    hipLaunchKernelGGL(gpu_traverse, dim3(1), dim3(1), 0, 0, ptr, steps);
    hipEventRecord(stop);
    hipEventSynchronize(stop);

    float millis = 0;
    if (PRINT_IF_ERROR(hipEventElapsedTime(&millis, start, stop))) {
      state.SkipWithError(NAME " failed to get elapsed time");
      break;
    }
    state.SetIterationTime(millis / 1000);
    state.counters["strides"] = steps;
  }
}

BENCHMARK(Comm_UM_Latency_GPUToGPU)->SMALL_ARGS()->UseManualTime();

#endif // CUDA_VERSION_MAJOR >= 8